from __future__ import print_function
import sys
import cv2
from random import randint
trackerTypes = ['BOOSTING', 'MIL', 'KCF','TLD', 'MEDIANFLOW', 'GOTURN', 'MOSSE', 'CSRT']
def createTrackerByName(trackerType):
    # Create a tracker based on tracker name
    if trackerType == trackerTypes[0]:
        tracker = cv2.TrackerBoosting_create()
    elif trackerType == trackerTypes[1]: 
        tracker = cv2.TrackerMIL_create()
    elif trackerType == trackerTypes[2]:
        tracker = cv2.TrackerKCF_create()
    elif trackerType == trackerTypes[3]:
        tracker = cv2.TrackerTLD_create()
    elif trackerType == trackerTypes[4]:
        tracker = cv2.TrackerMedianFlow_create()
    elif trackerType == trackerTypes[5]:
        tracker = cv2.TrackerGOTURN_create()
    elif trackerType == trackerTypes[6]:
        tracker = cv2.TrackerMOSSE_create()
    elif trackerType == trackerTypes[7]:
        tracker = cv2.TrackerCSRT_create()
    else:
        tracker = None
        print('错误的跟踪器的名字')
        print('我们有以下几种跟踪器:')
        for t in trackerTypes:
            print(t)
    return tracker
# 设置路径
videoPath = "1.mp4"
# 创建一个caputure类
cap=cv2.VideoCapture(1)
# 读取初帧
ret,frame = cap.read()
# 判断是否读取成功
if not ret:
    print('读取失败')
    sys.exit(1)
bboxes = []
colors = [] 
while 1:
    bbox = cv2.selectROI('MultiTracker', frame)
    bboxes.append(bbox)
    colors.append((randint(0, 255), randint(0, 255), randint(0, 255)))
    print("按下q键退出，按下其他键继续画下一个框")
    if cv2.waitKey(0) & 0xFF==ord('q'):
        break
print('选取的边框为：{}'.format(bboxes))
# 选择单个的跟踪器类型
trackerType = "CSRT"   
# 实例化
multiTracker = cv2.MultiTracker_create()
# 初始化跟踪器
for bbox in bboxes:
    multiTracker.add(createTrackerByName(trackerType), frame, bbox)
# 开始对后续帧进行跟踪
while cap.isOpened():
    ret,frame =cap.read()
    if not ret:
        break
    # 更新位置
    ret,boxes = multiTracker.update(frame)
    #显示是否出错
    print(ret)
    # 画出边界框
    for i,newbox in enumerate(boxes):
        p1=(int(newbox[0]), int(newbox[1]))
        p2=(int(newbox[0] + newbox[2]), int(newbox[1] + newbox[3]))
        cv2.rectangle(frame, p1, p2, colors[i], 2, 1)
    # 显示画面
    cv2.imshow('MultiTracker', frame)
    # 按下q退出
    if cv2.waitKey(1) & 0xFF == ord('q'): 
        break